package raft

import (
	"context"
	"fmt"
	"go.uber.org/zap"
	"google.golang.org/protobuf/proto"
	"net"
	"time"
	"xin/raft/pb"
)

const (
	connBufferSize = 256 * 1024

	RpcRequestVote uint8 = iota
	RpcRequestVoteResp
	RpcAppendEntries
	RpcAppendEntriesResp
)

type TcpTransport struct {
	port       int
	consumerCh chan *RPCReq
	listen     net.Listener
	connPool   map[string]net.Conn // 网络连接池
	logger     *zap.Logger
}

func NewTcpTransport(port int, logger *zap.Logger) *TcpTransport {
	return &TcpTransport{
		port:       port,
		consumerCh: make(chan *RPCReq, 1),
		connPool:   make(map[string]net.Conn),
		logger:     logger,
	}
}

func (t *TcpTransport) accept() {
	for {
		conn, err := t.listen.Accept()
		if err != nil {
			// TODO 等待一段时间重新接收。
			time.Sleep(time.Millisecond)
			continue
		}

		// 开启一个协程处理网络请求
		go t.processConn(context.TODO(), conn) // TODO context Shutdown 时使用
	}
}

func (t *TcpTransport) Init() error {
	listen, err := net.Listen("tcp", fmt.Sprintf(":%d", t.port))
	if err != nil {
		return err
	}

	t.listen = listen

	go t.accept()

	return nil
}

func (t *TcpTransport) Shutdown() error {
	return t.listen.Close()
}

// process 处理网络请求
// 接收其他节点发送过来的信息
// 例：领导节点发送来的心跳，日志复制。候选节点发送过来的选举请求等。
func (t *TcpTransport) processConn(ctx context.Context, conn net.Conn) {
	defer conn.Close()

	for {
		select {
		case <-ctx.Done():
			t.logger.Warn("tcp transport shutting down")
			return
		default:
		}

		rpcType, buff, err := t.decode(conn)
		if err != nil {
			return
		}

		var respChan = make(chan RPCResp, 1)

		var rpcReq = &RPCReq{
			RespChan: respChan,
		}

		switch rpcType {
		case RpcRequestVote: // 收到请求投票请求
			req := &pb.RequestVote{}
			err = proto.Unmarshal(buff, req)
			if err != nil {
				t.logger.Error("unmarshal rpc request vote", zap.Error(err))
				return
			}

			// t.logger.Debug(">>> 收到请求投票网络请求", zap.Any("vote", req))

			rpcReq.Command = req

			t.consumerCh <- rpcReq

		case RpcAppendEntries:
			// t.logger.Debug(">>> 收到追加日志请求")

			req := &pb.AppendEntries{}
			err = proto.Unmarshal(buff, req)
			if err != nil {
				t.logger.Error("unmarshal rpc  append entries", zap.Error(err))
				return
			}

			rpcReq.Command = req
			t.consumerCh <- rpcReq

		default:
			t.logger.Error("Unknown rpc type", zap.Uint8("type", rpcType))
			return
		}

		// 等待算法层返回结果
		select {
		case resp := <-respChan:
			switch result := resp.Resp.(type) {
			case *pb.RequestVoteResp:
				// t.logger.Debug(">>> 收到算法层，回复结果", zap.Uint64("term", result.Term), zap.Bool("VoteGranted", result.VoteGranted))

				bytes, err := t.encode(RpcRequestVoteResp, result)
				if err != nil {
					t.logger.Error("marshal rpc response vote", zap.Error(err))
					return
				}

				// t.logger.Debug(">>> 开始写入网络")
				_, err = conn.Write(bytes)
				if err != nil {
					t.logger.Error("write rpc response vote", zap.Error(err))
					return
				}
			case *pb.AppendEntriesResp:
				// t.logger.Debug(">>> 收到算法层，回复结果 AppendEntriesResp")

				// 收到算法层，日志处理回复
				bytes, err := t.encode(RpcAppendEntriesResp, result)
				if err != nil {
					t.logger.Error("marshal rpc response append entries", zap.Error(err))
					return
				}

				_, err = conn.Write(bytes)
				if err != nil {
					t.logger.Error("write rpc response append entries", zap.Error(err))
					return
				}
			}
		}
	}
}

// encode 编码 将 Protobuf 结构体序列化为 字节切片
// 第一个字节为消息类型
// 第二到五个字节为消息长度
// 剩下的为消息具体内容
func (t *TcpTransport) encode(rpcType uint8, msg proto.Message) ([]byte, error) {
	bytes, err := proto.Marshal(msg)
	if err != nil {
		t.logger.Error("marshal proto msg err", zap.Error(err))
		return nil, err
	}

	l := uint32(len(bytes)) // 消息长度
	packet := make([]byte, 5+len(bytes))
	packet[0] = rpcType
	packet[1] = byte(l >> 24)
	packet[2] = byte(l >> 16)
	packet[3] = byte(l >> 8)
	packet[4] = byte(l)

	copy(packet[5:], bytes)

	return packet, nil
}

// decode 解码 从 conn 读取字节信息
func (t *TcpTransport) decode(conn net.Conn) (byte, []byte, error) {
	// 读取第一个字节，表示消息类型
	rpcType := make([]byte, 1)
	_, err := conn.Read(rpcType)
	if err != nil {
		return 0, nil, err
	}

	// 读取 4 个字节，表示消息长度
	lenBuf := make([]byte, 4)
	_, err = conn.Read(lenBuf)
	if err != nil {
		return 0, nil, err
	}

	l := uint32(lenBuf[0])<<24 | uint32(lenBuf[1])<<16 | uint32(lenBuf[2])<<8 | uint32(lenBuf[3])

	buff := make([]byte, l)
	_, err = conn.Read(buff)
	if err != nil {
		return 0, nil, err
	}

	return rpcType[0], buff, nil
}

func (t *TcpTransport) getConn(endpoint Endpoint) (net.Conn, error) {
	if conn, ok := t.connPool[endpoint.String()]; ok {
		return conn, nil
	}

	conn, err := net.Dial("tcp", endpoint.String())
	if err != nil {
		t.logger.Error("dial tcp endpoint error", zap.Error(err))
		return nil, err
	}

	t.connPool[endpoint.String()] = conn

	return conn, nil
}

// sendRpc 发送 rpc request
func (t *TcpTransport) sendRpc(req proto.Message, rpcType uint8, endpoint Endpoint, resp proto.Message) error {
	conn, err := t.getConn(endpoint)
	if err != nil {
		t.logger.Error("获取 TCP 连接错误", zap.Error(err))
		return err
	}

	packet, err := t.encode(rpcType, req)
	if err != nil {
		t.logger.Error("请求参数编码错误", zap.Error(err))
		return err
	}

	if _, err := conn.Write(packet); err != nil {
		t.logger.Error("写入网络数据包错误", zap.Error(err))
		return err
	}

	_, bytes, err := t.decode(conn)
	if err != nil {
		t.logger.Error("解码请求返回结果错误", zap.Error(err))
		return err
	}

	err = proto.Unmarshal(bytes, resp)
	if err != nil {
		t.logger.Error("unmarshal rpc response", zap.Error(err))
		return err
	}

	// t.logger.Debug(">>> rpc response", zap.Any("resp", resp))
	return nil
}

func (t *TcpTransport) RequestVote(req *pb.RequestVote, endpoint Endpoint, resp *pb.RequestVoteResp) error {
	// t.logger.Debug(">>> 发送请求投票", zap.String("endpoint", endpoint.String()), zap.Any("req", req))
	return t.sendRpc(req, RpcRequestVote, endpoint, resp)
}

func (t *TcpTransport) AppendEntries(req *pb.AppendEntries, endpoint Endpoint, resp *pb.AppendEntriesResp) error {
	// t.logger.Debug(">>> 发送追加日志", zap.String("endpoint", endpoint.String()), zap.Any("req", req))
	return t.sendRpc(req, RpcAppendEntries, endpoint, resp)
}

func (t *TcpTransport) Consumer() chan *RPCReq {
	return t.consumerCh
}
